package onnxruntime

import (
	"errors"
	"fmt"
	"os"
	"unsafe"
)

// #cgo CFLAGS: -O2 -g
//
// #include "onnxruntime_wrapper.h"
import "C"

type TensorValue struct {
	Value interface{}
	Shape []int64
}

type Session struct {
	ortSession  *C.OrtSession
	inputNames  []*C.char
	outputNames []*C.char
}

func NewSession(onnxFilePath string, inputNames []string, outputNames []string,
	options *SessionOptions) (*Session, error) {

	fileContent, e := os.ReadFile(onnxFilePath)
	if e != nil {
		return nil, fmt.Errorf("error reading %s: %w", onnxFilePath, e)
	}
	ortSession, e := createCSession(fileContent, options)
	if e != nil {
		return nil, fmt.Errorf("error creating C session: %w", e)
	}

	cInputNames := make([]*C.char, len(inputNames))
	cOutputNames := make([]*C.char, len(outputNames))
	for i, v := range inputNames {
		cInputNames[i] = C.CString(v)
	}
	for i, v := range outputNames {
		cOutputNames[i] = C.CString(v)
	}

	// We don't use the input and output list of OrtValues with these.
	s := &Session{
		ortSession:  ortSession,
		inputNames:  cInputNames,
		outputNames: cOutputNames,
	}
	return s, nil
}

func (s *Session) Destroy() error {
	if s.ortSession != nil {
		C.ReleaseOrtSession(s.ortSession)
		s.ortSession = nil
	}
	for i := range s.inputNames {
		C.free(unsafe.Pointer(s.inputNames[i]))
	}
	s.inputNames = nil
	for i := range s.outputNames {
		C.free(unsafe.Pointer(s.outputNames[i]))
	}
	s.outputNames = nil
	return nil
}

// Runs the session, updating the contents of the output tensors on success.
func (s *Session) Predict(inputs []ArbitraryTensor, outputs []ArbitraryTensor) error {

	if !IsInitialized() {
		return ErrorNotInitialized
	}

	if len(inputs) == 0 {
		return errors.New("no inputs were provided")
	}
	if len(inputs) != len(s.inputNames) {
		return fmt.Errorf("got %d input tensors, but %d input names",
			len(inputs), len(s.inputNames))
	}

	inputValues := make([]*C.OrtValue, len(inputs))
	for i, v := range inputs {
		inputValues[i] = v.GetInternals().OrtValue
	}

	output := make([]*C.OrtValue, len(outputs))
	for i, v := range outputs {
		output[i] = v.GetInternals().OrtValue
	}

	status := C.RunOrtSession(s.ortSession, &inputValues[0], &s.inputNames[0],
		C.int(len(inputs)), &output[0], &s.outputNames[0],
		C.int(1))
	if status != nil {
		return fmt.Errorf("error running network: %w", statusToError(status))
	}

	return nil
}

func createCSession(onnxData []byte, options *SessionOptions) (*C.OrtSession, error) {
	if !IsInitialized() {
		return nil, ErrorNotInitialized
	}
	if len(onnxData) == 0 {
		return nil, errors.New("missing ONNX data")
	}
	var ortSession *C.OrtSession
	var ortSessionOptions *C.OrtSessionOptions
	if options != nil {
		ortSessionOptions = options.o
	}
	status := C.CreateSession(unsafe.Pointer(&(onnxData[0])),
		C.size_t(len(onnxData)), ortEnv, &ortSession, ortSessionOptions)
	if status != nil {
		return nil, statusToError(status)
	}
	return ortSession, nil
}

func createTensorFromOrtValue(v *C.OrtValue) (ArbitraryTensor, error) {
	var pInfo *C.OrtTensorTypeAndShapeInfo
	status := C.GetTensorTypeAndShape(v, &pInfo)
	if status != nil {
		return nil, fmt.Errorf("error getting type and shape: %w", statusToError(status))
	}
	var dimCount C.size_t
	status = C.GetDimensionsCount(pInfo, &dimCount)
	if status != nil {
		return nil, fmt.Errorf("error getting dimensions count: %w", statusToError(status))
	}
	shape := make(Shape, dimCount)
	status = C.GetDimensions(pInfo, (*C.int64_t)(&shape[0]), dimCount)
	if status != nil {
		return nil, fmt.Errorf("error getting dimensions: %w", statusToError(status))
	}
	var tensorElementType C.ONNXTensorElementDataType
	status = C.GetTensorElementType(pInfo, (*uint32)(&tensorElementType))
	if status != nil {
		return nil, fmt.Errorf("error getting element type: %w", statusToError(status))
	}
	C.ReleaseTensorTypeAndShapeInfo(pInfo)
	var tensorData unsafe.Pointer
	status = C.GetTensorMutableData(v, &tensorData)
	if status != nil {
		return nil, fmt.Errorf("error getting tensor mutable data: %w", statusToError(status))
	}

	switch tensorType := TensorElementDataType(tensorElementType); tensorType {
	case TensorElementDataTypeFloat:
		return createTensorWithCData[float32](shape, tensorData)
	case TensorElementDataTypeUint8:
		return createTensorWithCData[uint8](shape, tensorData)
	case TensorElementDataTypeInt8:
		return createTensorWithCData[int8](shape, tensorData)
	case TensorElementDataTypeUint16:
		return createTensorWithCData[uint16](shape, tensorData)
	case TensorElementDataTypeInt16:
		return createTensorWithCData[int16](shape, tensorData)
	case TensorElementDataTypeInt32:
		return createTensorWithCData[int32](shape, tensorData)
	case TensorElementDataTypeInt64:
		return createTensorWithCData[int64](shape, tensorData)
	case TensorElementDataTypeDouble:
		return createTensorWithCData[float64](shape, tensorData)
	case TensorElementDataTypeUint32:
		return createTensorWithCData[uint32](shape, tensorData)
	case TensorElementDataTypeUint64:
		return createTensorWithCData[uint64](shape, tensorData)
	default:
		totalSize := shape.FlattenedSize()
		actualData := unsafe.Slice((*byte)(tensorData), totalSize)
		return NewCustomDataTensor(shape, actualData, tensorType)
	}
}

func createTensorWithCData[T TensorData](shape Shape, data unsafe.Pointer) (*Tensor[T], error) {
	totalSize := shape.FlattenedSize()
	actualData := unsafe.Slice((*T)(data), totalSize)
	return NewTensor[T](shape, actualData)
}
